import time
import m3u8
import requests
import os
import subprocess
import urllib.parse
from multiprocessing import Pool

def download_segment(segment_info):
    ts_url, output_folder, m3u8_url, cookies, proxies, index = segment_info
    if not ts_url.startswith('http'):
        ts_url = urllib.parse.urljoin(m3u8_url, ts_url)
    local_ts_path = os.path.join(output_folder, os.path.basename(ts_url).split('?')[0])
    proxy_config = {'http': proxies, 'https': proxies} if proxies else {}

    try:
        response = requests.get(ts_url, stream=True, cookies=cookies, proxies=proxy_config)
        response.raise_for_status()
        with open(local_ts_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return local_ts_path, index
    except Exception as e:
        print(f"Error downloading segment: {e}")
        return None, index

def merge_videos(playlist_path, output_path):
    command = [
        './ffmpeg',
        '-i', playlist_path,
        '-c', 'copy',
        output_path
    ]
    try:
        subprocess.run(command, check=True)
        print("Video merging completed!")
        # 合并完成后删除所有TS文件和M3U8文件
        cleanup_files(playlist_path)
    except subprocess.CalledProcessError as e:
        print(f"Error merging video: {e}")

def fix_m3u8_content(m3u8_text):
    lines = m3u8_text.splitlines(True)
    fixed_lines = []
    seen_ts_files = set()
    media_sequence_found = False

    for line in lines:
        if '#EXT-X-MEDIA-SEQUENCE' in line:
            if not media_sequence_found:
                media_sequence_found = True
                fixed_lines.append(line)
        elif line.strip().endswith('.ts'):
            if line not in seen_ts_files:
                seen_ts_files.add(line)
                fixed_lines.append(line)
        else:
            fixed_lines.append(line)

    return ''.join(fixed_lines)

def save_playlist(output_folder, local_m3u8):
    # 假设这个函数负责保存你的播放列表到 'local_playlist.m3u8'
    # 具体保存逻辑根据实际情况编写
    with open(os.path.join(output_folder, 'local_playlist.m3u8'), 'w') as f:
        m3u8_text = local_m3u8.dumps().replace('#EXTM3U\n', '#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-TARGETDURATION:4\n#EXT-X-MEDIA-SEQUENCE:0\n')+'\n'
        finall_m3u8 = fix_m3u8_content(m3u8_text)
        if '#EXT-X-ENDLIST' not in finall_m3u8:
            finall_m3u8 += '#EXT-X-ENDLIST\n'
        f.write(finall_m3u8)
    merge_videos(f'{output_folder}/local_playlist.m3u8', f'{output_folder}/{int(time.time())}.mp4')

def cleanup_files(playlist_path):
    """删除所有TS文件和M3U8文件"""
    directory = os.path.dirname(playlist_path)

    # 删除M3U8文件
    if os.path.exists(playlist_path):
        os.remove(playlist_path)
        print(f"Deleted: {playlist_path}")

    # 删除对应目录中的所有TS文件
    for file_name in os.listdir(directory):
        if file_name.endswith('.ts'):
            ts_file_path = os.path.join(directory, file_name)
            os.remove(ts_file_path)
            print(f"Deleted: {ts_file_path}")

def download_ts_files(m3u8_url, output_folder, progress_callback, running_flag, cookies='', proxies=None):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    proxy_config = {'http': proxies, 'https': proxies} if proxies else {}
    local_m3u8 = m3u8.M3U8()

    while running_flag.value:
        try:
            response = requests.get(m3u8_url, cookies=cookies, proxies=proxy_config)
            response.raise_for_status()

            m3u8_parsed = m3u8.loads(response.text)
            is_live = '#EXT-X-ENDLIST' not in response.text

            segments_info = [
                (segment.uri, output_folder, m3u8_url, cookies, proxies, i)
                for i, segment in enumerate(m3u8_parsed.segments)
            ]

            with Pool(processes=4) as pool:
                for local_ts_path, index in pool.imap_unordered(download_segment, segments_info):
                    if local_ts_path:
                        segment = m3u8_parsed.segments[index]
                        local_m3u8.add_segment(m3u8.Segment(uri=os.path.basename(local_ts_path), duration=segment.duration))
                        progress_callback(index + 1, len(m3u8_parsed.segments))

            if not is_live:
                break

            time.sleep(1)

        except Exception as e:
            print(f"Error during download: {e}")
            break

    save_playlist(output_folder, local_m3u8)

if __name__ == '__main__':
    download_ts_files(m3u8_url='', output_folder='./output')  # type: ignore
